#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多智能体系统实现

采用协调者模式（Coordinator Pattern）：
- 协调者 Agent：负责识别用户意图并分发给专家
- 专家 Agent：各自负责特定领域的任务处理
"""

from llama_index.core.agent import ReActAgent, ReActChatFormatter
from llama_index.core.agent.react.prompts import REACT_CHAT_SYSTEM_HEADER
from llama_index.core.tools import FunctionTool
from llama_index.core.memory import ChatMemoryBuffer
from biz.config.settings import openLLM, cache
from biz.utils.logger import logger


class MultiAgentSystem:
    """多智能体系统 - 协调者模式"""
    
    def __init__(self):
        """初始化多智能体系统"""
        # 创建专家智能体
        self.sql_agent = None
        self.command_agent = None
        self.knowledge_agent = None
        self.router_agent = None
        self.diagnosis_agent = None
        
        # 创建协调者智能体
        self.coordinator = None
    
    def _create_sql_agent(self, session_id: str, memory: ChatMemoryBuffer):
        """创建SQL分析专家"""
        from biz.core.functions import generate_sql_core
        
        def sql_analysis_tool_fn(question: str) -> str:
            """SQL分析工具"""
            try:
                id, sql = generate_sql_core(question, session_id)
                return sql
            except Exception as e:
                logger.error(f"SQL分析失败: {e}")
                return f"SQL生成失败: {str(e)}"
        
        sql_tool = FunctionTool.from_defaults(
            fn=sql_analysis_tool_fn,
            name="sql_analysis",
            description="执行数据查询和分析任务，生成SQL语句",
            return_direct=True
        )
        
        return ReActAgent.from_tools(
            tools=[sql_tool],
            llm=openLLM,
            memory=memory,
            system_prompt="""你是数据分析专家。

职责：
- 处理用户的数据查询、统计分析、报表生成需求
- 使用 sql_analysis 工具生成SQL查询语句
- 直接返回SQL结果，不要添加额外说明

工作流程：
1. 理解用户的数据查询需求
2. 调用 sql_analysis 工具生成SQL
3. 返回生成的SQL语句

注意：
- 只处理数据查询相关的任务
- 返回格式简洁明了
- 如果遇到问题，说明具体原因""",
            max_iterations=1,
            verbose=False,
        )
    
    def _create_command_agent(self, session_id: str, memory: ChatMemoryBuffer):
        """创建设备控制专家"""
        from biz.core.functions import generate_command_core
        
        def command_control_tool_fn(question: str) -> str:
            """设备控制工具"""
            try:
                uuid, cmd_type, message = generate_command_core(question, session_id, "command")
                return message
            except Exception as e:
                logger.error(f"设备控制失败: {e}")
                return f"控制命令执行失败: {str(e)}"
        
        command_tool = FunctionTool.from_defaults(
            fn=command_control_tool_fn,
            name="device_control",
            description="执行设备控制命令，如开关灯、调光等操作"
        )
        
        return ReActAgent.from_tools(
            tools=[command_tool],
            llm=openLLM,
            memory=memory,
            system_prompt="""你是设备控制专家。

职责：
- 处理用户的设备控制需求（开关灯、调光等）
- 使用 device_control 工具执行控制命令
- 直接返回执行结果，不要添加额外说明

工作流程：
1. 理解用户的控制需求
2. 调用 device_control 工具执行命令
3. 返回执行结果

注意：
- 只处理设备控制相关的任务
- 不支持下发方案和计划功能
- 返回格式简洁明了""",
            max_iterations=1,
            verbose=False
        )
    
    def _create_knowledge_agent(self, memory: ChatMemoryBuffer):
        """创建知识库专家"""
        
        def knowledge_query_tool_fn(question: str) -> str:
            """适用于用户希望获取照明方面的知识或文档信息的情况。
            参数：
                question (str): 要查询的问题描述。
            """
            return question
        
        knowledge_tool = FunctionTool.from_defaults(
            fn=knowledge_query_tool_fn,
            name="knowledge_query",
            description="查询照明相关的知识、文档、FAQ",
            return_direct=True
        )
        
        return ReActAgent.from_tools(
            tools=[knowledge_tool],
            llm=openLLM,
            memory=memory,
            system_prompt="""你是知识库专家。

职责：
- 回答用户关于照明系统的知识咨询
- 提供操作指南、文档说明、常见问题解答
- 使用 knowledge_query 工具查询知识库

工作流程：
1. 理解用户的知识咨询需求
2. 调用 knowledge_query 工具查询
3. 返回知识内容

注意：
- 只处理知识咨询相关的任务
- 回答要准确、专业
- 返回格式清晰易懂""",
            max_iterations=1,
            verbose=False
        )
    
    def _create_router_agent(self, session_id: str, memory: ChatMemoryBuffer):
        """创建路由/菜单专家"""
        from biz.core.functions import generate_command_core
        
        def router_tool_fn(question: str) -> str:
            """路由工具"""
            try:
                uuid, cmd_type, message = generate_command_core(question, session_id, "router")
                return message
            except Exception as e:
                logger.error(f"路由处理失败: {e}")
                return f"菜单导航失败: {str(e)}"
        
        router_tool = FunctionTool.from_defaults(
            fn=router_tool_fn,
            name="menu_navigation",
            description="处理系统菜单导航和功能跳转",
            return_direct=True
        )
        
        return ReActAgent.from_tools(
            tools=[router_tool],
            llm=openLLM,
            memory=memory,
            system_prompt="""你是菜单导航专家。

职责：
- 处理用户的菜单查看和功能跳转需求
- 使用 menu_navigation 工具进行导航
- 直接返回导航结果

工作流程：
1. 理解用户的导航需求
2. 调用 menu_navigation 工具
3. 返回导航结果

注意：
- 只处理菜单导航相关的任务
- 返回格式简洁明了""",
            max_iterations=1,
            verbose=False
        )
    
    def _create_diagnosis_agent(self, memory: ChatMemoryBuffer):
        """创建诊断分析专家"""
        from biz.core.functions import generate_diagnosis_core
        
        def diagnosis_tool_fn(question: str) -> str:
            """诊断分析工具"""
            try:
                id, diagnosis_result = generate_diagnosis_core(question)
                return diagnosis_result
            except Exception as e:
                logger.error(f"诊断分析失败: {e}")
                return f"诊断分析失败: {str(e)}"
        
        diagnosis_tool = FunctionTool.from_defaults(
            fn=diagnosis_tool_fn,
            name="fault_diagnosis",
            description="对设备故障进行深度分析诊断，生成诊断报告",
            return_direct=True
        )
        
        return ReActAgent.from_tools(
            tools=[diagnosis_tool],
            llm=openLLM,
            memory=memory,
            system_prompt="""你是故障诊断专家。

职责：
- 对设备故障进行深度分析诊断
- 使用 fault_diagnosis 工具生成诊断报告
- 返回详细的诊断结果

工作流程：
1. 理解用户的诊断需求
2. 调用 fault_diagnosis 工具进行分析
3. 返回诊断报告

注意：
- 诊断过程可能需要较长时间
- 会生成详细的诊断报告
- 返回完整的诊断结果""",
            max_iterations=1,
            verbose=False
        )
    
    def _create_coordinator(self, session_id: str, memory: ChatMemoryBuffer):
        """创建协调者 - 负责识别意图并分发给专家"""
        
        # 延迟创建专家（避免循环依赖）
        def delegate_to_sql_expert(question: str) -> str:
            """委托给SQL分析专家"""
            logger.info(f"协调者：将任务分发给SQL分析专家 - {question[:50]}")
            if self.sql_agent is None:
                self.sql_agent = self._create_sql_agent(session_id, memory)
            response = self.sql_agent.chat(question)
            return str(response)
        
        def delegate_to_command_expert(question: str) -> str:
            """委托给设备控制专家"""
            logger.info(f"协调者：将任务分发给设备控制专家 - {question[:50]}")
            if self.command_agent is None:
                self.command_agent = self._create_command_agent(session_id, memory)
            response = self.command_agent.chat(question)
            return str(response)
        
        def delegate_to_knowledge_expert(question: str) -> str:
            """委托给知识库专家"""
            logger.info(f"协调者：将任务分发给知识库专家 - {question[:50]}")
            if self.knowledge_agent is None:
                self.knowledge_agent = self._create_knowledge_agent(memory)
            response = self.knowledge_agent.chat(question)
            return str(response)
        
        def delegate_to_router_expert(question: str) -> str:
            """委托给路由专家"""
            logger.info(f"协调者：将任务分发给路由专家 - {question[:50]}")
            if self.router_agent is None:
                self.router_agent = self._create_router_agent(session_id, memory)
            response = self.router_agent.chat(question)
            return str(response)
        
        def delegate_to_diagnosis_expert(question: str) -> str:
            """委托给诊断专家"""
            logger.info(f"协调者：将任务分发给诊断专家 - {question[:50]}")
            if self.diagnosis_agent is None:
                self.diagnosis_agent = self._create_diagnosis_agent(memory)
            response = self.diagnosis_agent.chat(question)
            return str(response)
        
        def ask_user_for_clarification(clarification_question: str) -> str:
            """向用户询问以澄清意图"""
            logger.info(f"协调者：需要用户澄清意图 - {clarification_question[:50]}")
            return clarification_question
        
        def end_conversation_fn(feedback: str) -> str:
            """结束对话"""
            logger.info(f"协调者：结束对话 - {feedback[:50]}")
            return feedback
        
        # 创建工具列表
        tools = [
            FunctionTool.from_defaults(
                fn=delegate_to_sql_expert,
                name="sql_expert",
                description="数据分析专家：处理数据查询、统计分析、报表生成等任务。适用于用户需要查询数据、分析统计信息的场景。",
                return_direct=True
            ),
            FunctionTool.from_defaults(
                fn=delegate_to_command_expert,
                name="control_expert",
                description="设备控制专家：处理设备控制、参数调整、操作执行等任务。适用于用户需要控制设备（开关灯、调光）的场景。",
                return_direct=True
            ),
            FunctionTool.from_defaults(
                fn=delegate_to_knowledge_expert,
                name="knowledge_expert",
                description="知识库专家：处理产品咨询、使用方法、常见问题等任务。适用于用户咨询照明知识、操作指南的场景。",
                return_direct=True
            ),
            FunctionTool.from_defaults(
                fn=delegate_to_router_expert,
                name="router_expert",
                description="菜单导航专家：处理系统菜单查看和功能跳转。适用于用户想要打开某个功能页面的场景。",
                return_direct=True
            ),
            FunctionTool.from_defaults(
                fn=delegate_to_diagnosis_expert,
                name="diagnosis_expert",
                description="故障诊断专家：对设备故障进行深度分析诊断。适用于用户需要诊断设备故障、分析问题原因的场景。",
                return_direct=True
            ),
            FunctionTool.from_defaults(
                fn=ask_user_for_clarification,
                name="ask_clarification",
                description="向用户询问：当用户意图不明确或需要更多信息时使用。生成一个问题来澄清用户的真实需求。",
                return_direct=True
            ),
            FunctionTool.from_defaults(
                fn=end_conversation_fn,
                name="end_conversation",
                description="结束对话：当用户表示结束对话或问题已解决时使用。",
                return_direct=True
            )
        ]
        
        # 协调者的系统提示词
        coordinator_prompt = f"""你是智能照明系统的协调者，负责理解用户意图并分发给合适的专家处理。

当前会话ID: {session_id}

## 可用专家

1. **sql_expert** - 数据分析专家
   - 处理：数据查询、统计分析、报表生成
   - 示例：查询昨天的能耗、统计故障数量、导出设备状态

2. **control_expert** - 设备控制专家
   - 处理：设备控制、开关灯、调光
   - 示例：打开路灯、调整亮度、关闭设备

3. **knowledge_expert** - 知识库专家
   - 处理：知识咨询、操作指南、FAQ
   - 示例：如何安装路灯、什么是智能照明、故障排查方法

4. **router_expert** - 菜单导航专家
   - 处理：功能跳转、菜单导航
   - 示例：打开灯杆管理、跳转到设备列表

5. **diagnosis_expert** - 故障诊断专家
   - 处理：设备故障深度诊断、问题分析
   - 示例：诊断设备为什么掉线、分析故障原因

6. **ask_clarification** - 询问用户
   - 使用场景：用户意图不明确、需要更多信息
   - 生成一个问题来澄清用户需求

7. **end_conversation** - 结束对话
   - 使用场景：用户表示结束、问题已解决

## 工作流程

1. **分析用户意图**：理解用户的真实需求
2. **选择专家**：根据意图选择最合适的专家
3. **委托处理**：调用专家工具处理任务
4. **返回结果**：直接返回专家的回复，不要添加额外说明

## 重要规则

- 必须选择一个工具来处理用户请求
- 如果意图不明确，使用 ask_clarification 询问用户
- 直接返回专家的回复，不要说"我已经..."等说明
- 每次只调用一个专家
- 优先使用最匹配的专家"""
        
        return ReActAgent.from_tools(
            tools=tools,
            llm=openLLM,
            memory=memory,
            react_chat_formatter=ReActChatFormatter(
                system_header=coordinator_prompt + "\n" + REACT_CHAT_SYSTEM_HEADER + "\n 你回复的内容中必须包含**Action**和**Action Input**！"
            ),
            max_iterations=5,
            verbose=False
        )
    
    def process(self, question: str, id: str = None, max_attempts: int = 5) -> tuple:
        """
        处理用户问题（识别 + 分发）
        
        Args:
            question: 用户问题
            id: 会话ID
            max_attempts: 最大重试次数
        
        Returns:
            tuple: (type, response)
                - type: 'sql' | 'command' | 'knowledge' | 'router' | 'diagnosis' | 'other' | 'complete'
                - response: 处理结果
        """
        try:
            # 获取或创建记忆
            if cache.get(id=id, field='memory') is not None:
                memory = cache.get(id=id, field='memory')
            else:
                memory = ChatMemoryBuffer.from_defaults(token_limit=500)
                cache.set(id=id, field='memory', value=memory)
            
            logger.info(f"多智能体系统接收问题: {question[:100]}")
            
            # 创建协调者（如果还没创建）
            if self.coordinator is None:
                self.coordinator = self._create_coordinator(id, memory)
            
            # 重试机制（只在格式错误时重试）
            attempt_count = 0
            task_type = 'other'
            
            while attempt_count < max_attempts:
                attempt_count += 1
                
                try:
                    # 协调者处理问题
                    response = self.coordinator.chat(question)
                    
                    # 检查是否有工具调用
                    if not response.sources or not response.sources[-1].tool_name:
                        # 只有在格式错误时才重试
                        if attempt_count < max_attempts:
                            question = '回复信息格式验证失败，回复内容中必须包含**Action**和**Action Input**！'
                            logger.warning(f"格式验证失败，重试 {attempt_count}/{max_attempts}")
                            continue
                        else:
                            # 达到最大重试次数
                            return 'other', "抱歉，系统似乎遇到了一点小问题，建议您稍后重新尝试。"
                    
                    # 根据调用的工具判断类型
                    tool_name = response.sources[-1].tool_name
                    
                    type_mapping = {
                        'sql_expert': 'sql',
                        'control_expert': 'command',
                        'knowledge_expert': 'knowledge',
                        'router_expert': 'router',
                        'diagnosis_expert': 'diagnosis',
                        'ask_clarification': 'other',
                        'end_conversation': 'complete'
                    }
                    
                    task_type = type_mapping.get(tool_name, 'other')
                    
                    logger.info(f"任务类型: {task_type}, 使用工具: {tool_name}")
                    
                    # 成功调用工具，直接返回（不再重试）
                    return task_type, str(response)
                    
                except Exception as e:
                    logger.error(f"处理失败，尝试 {attempt_count}/{max_attempts}: {e}")
                    if attempt_count < max_attempts:
                        question = '回复信息格式验证失败，回复内容中必须包含**Action**和**Action Input**！'
                        continue
                    else:
                        return 'other', f"系统错误: {str(e)}"
            
            # 不应该到达这里
            return task_type, "抱歉，系统似乎遇到了一点小问题，建议您稍后重新尝试。"
            
        except Exception as e:
            logger.error(f"多智能体处理失败: {e}", exc_info=True)
            return 'other', "抱歉，系统遇到问题，请稍后重试"


# 全局实例（单例模式）
_multi_agent_system = None

def get_multi_agent_system():
    """获取多智能体系统单例"""
    global _multi_agent_system
    if _multi_agent_system is None:
        _multi_agent_system = MultiAgentSystem()
    return _multi_agent_system
